import os
from PIL import Image
from concurrent.futures import ThreadPoolExecutor
import multiprocessing

def compress_image(input_file, output_file, quality=95, log_callback=None):
    """压缩单个图像文件，针对PNG使用优化压缩"""
    try:
        with Image.open(input_file) as img:
            file_format = img.format
            
            # 针对PNG格式的特殊处理
            if file_format == 'PNG':
                # 转换为P模式（调色板模式）以减小文件大小
                if img.mode not in ['P', 'L']:
                    img = img.convert('P', palette=Image.Palette.ADAPTIVE, colors=256)
                
                # 使用更高效的压缩设置
                img.save(
                    output_file, 
                    format='PNG',
                    optimize=True,
                    compress_level=9  # 最高压缩级别
                )
            else:
                # 其他格式的压缩
                try:
                    img.save(output_file, 
                            format=file_format, 
                            optimize=True,
                            quality=quality)
                except OSError:
                    img.save(output_file, 
                            format=file_format, 
                            optimize=True, 
                            exif=b'',
                            quality=quality)
        
        # 计算压缩率
        orig_size = os.path.getsize(input_file)
        comp_size = os.path.getsize(output_file)
        ratio = (1 - comp_size / orig_size) * 100 if orig_size > 0 else 0
        
        msg = f"✓ 已压缩: {os.path.basename(output_file)} - 压缩率: {ratio:.1f}%"
        if log_callback: log_callback(msg)
    except Exception as e:
        msg = f"✗ 处理失败: {input_file} - 原因: {str(e)}"
        if log_callback: log_callback(msg)

def compress_images(input_path, output_path, extensions, quality, overwrite=False, thread_count=None, log_callback=None):
    """批量压缩图像文件（支持文件/文件夹）"""
    supported_formats = ['JPEG', 'PNG', 'BMP', 'GIF', 'TIFF']
    image_files = []
    
    if thread_count is None:
        thread_count = multiprocessing.cpu_count()
    
    # 确保输出目录存在
    if not os.path.exists(output_path):
        os.makedirs(output_path)
    
    # 判断输入类型
    if os.path.isdir(input_path):  # 文件夹模式
        if log_callback: log_callback(f"📂 处理文件夹: {input_path}")
        for root, _, files in os.walk(input_path):
            for file in files:
                input_file = os.path.join(root, file)
                _, ext = os.path.splitext(input_file)
                
                if ext.lower() in extensions:
                    try:
                        with Image.open(input_file) as img:
                            if img.format in supported_formats:
                                image_files.append(input_file)
                    except Exception as e:
                        msg = f"⚠ 跳过无效文件: {os.path.basename(input_file)} - {str(e)}"
                        if log_callback: log_callback(msg)
    
    elif os.path.isfile(input_path):  # 单文件模式
        if log_callback: log_callback(f"🖼 处理文件: {os.path.basename(input_path)}")
        _, ext = os.path.splitext(input_path)
        if ext.lower() in extensions:
            try:
                with Image.open(input_path) as img:
                    if img.format in supported_formats:
                        image_files.append(input_path)
            except Exception as e:
                msg = f"⚠ 跳过无效文件: {os.path.basename(input_path)} - {str(e)}"
                if log_callback: log_callback(msg)
    else:
        msg = f"❌ 错误：路径不存在 - {input_path}"
        if log_callback: log_callback(msg)
        return

    # 并行压缩处理
    if image_files:
        with ThreadPoolExecutor(max_workers=thread_count) as executor:
            for input_file in image_files:
                # 构建输出路径
                if os.path.isdir(input_path):
                    # 保持相对路径结构
                    rel_path = os.path.relpath(os.path.dirname(input_file), input_path)
                    output_dir = os.path.join(output_path, rel_path)
                    if not os.path.exists(output_dir):
                        os.makedirs(output_dir)
                    
                    base_name = os.path.basename(input_file)
                    output_file = os.path.join(output_dir, base_name)
                else:
                    # 单文件模式
                    base_name = os.path.basename(input_file)
                    output_file = os.path.join(output_path, base_name)
                
                # 处理文件名冲突（根据覆盖选项）
                base, ext = os.path.splitext(output_file)
                if overwrite:
                    # 覆盖模式：直接使用目标路径
                    final_output = output_file
                    if os.path.exists(final_output):
                        if log_callback: log_callback(f"⚠ 覆盖文件: {os.path.basename(final_output)}")
                else:
                    # 非覆盖模式：添加序号避免冲突
                    candidate = output_file
                    counter = 1
                    while os.path.exists(candidate):
                        candidate = f"{base}_{counter}{ext}"
                        counter += 1
                    final_output = candidate
                
                executor.submit(compress_image, input_file, final_output, quality, log_callback)
        
        if overwrite:
            if log_callback: log_callback(f"✅ 所有任务已完成！(覆盖模式)")
        else:
            if log_callback: log_callback("✅ 所有任务已完成！")
    else:
        if log_callback: log_callback("⛔ 未找到有效图像文件")